from math import log,inf
from itertools import product
from module import Matrix,Network,Node

class HiddenMarkovModel:
    
    def __init__(self,init_probs:   dict[float],
                 emission_probs:    dict[dict],
                 trans_probs:       dict[dict]) -> None:
        self.initstate_probs = init_probs
        self.emission_probs = Matrix(emission_probs)
        self.transition_probs = Matrix(trans_probs)
        self.states,self.symbols = self.emission_probs.keys()
        self.set_log()

    def set_log(self,base=2) -> None:
        self.base = base
        self.I_log = {state:log(self.initstate_probs[state],base)
                      for state in self.states}
        self.E_log = self.emission_probs.log(base)
        self.T_log = self.transition_probs.log(base)
    
    def view_parameters(self) -> None:
        print ("\nEmission probabilities")
        self.emission_probs.View()
        print ("Transition probabilities")
        self.transition_probs.View()
        print (f"Take log base of {self.base}")
        print ("log emission probabilities")
        self.E_log.View()
        print ("log transition probabilities")
        self.T_log.View()
    
    def joint_score(self, observed_seq, state_path) -> float:
        if not observed_seq:
            raise ValueError("Empty observed sequence")
        if len(state_path) != len(observed_seq):
            raise ValueError("Observed sequence and state path of different lengths")
        
        score = self.I_log[state_path[0]] + \
                self.E_log[state_path[0],observed_seq[0]]
        
        # for obs,sta_from,sta_to in zip(observed_seq[1:],state_path[:-1],state_path[1:]):
        #     score += self.transition_probs[sta_from,sta_to] + self.emission_probs[sta_to,obs]
        
        score += sum(
            self.T_log[sta_from,sta_to] + self.E_log[sta_to,obs]
            for obs,sta_from,sta_to in zip(observed_seq[1:],
                                           state_path[:-1],
                                           state_path[1:])
            )
        
        return score + len(state_path) * log(4,self.base)

    def viterbi(self,seq):
        N = Network(self.states,seq)
        for i in range(1,len(seq)):
            for state in self.states:
                N[i,state].set_value(self.E_log[state,seq[i]])
            for sta_from,sta_to in product(self.states,self.states):
                if self.T_log[sta_from,sta_to] == -inf:
                    continue
                N[i-1,sta_from].add_next(N[i,sta_to],self.T_log[sta_from,sta_to])
        
        StartNode = Node('',-1)
        StartNode.set_value(0)
        for state in self.states:
            N[0,state].set_value(self.E_log[state,seq[0]])
            StartNode.add_next(N[0,state],self.I_log[state])
        
        return N.all_path(StartNode)
        

def test_HMM():
    initial_probs = {"5": 0.8,"M": 0.15,"3": 0.05}    
    emission_probs = {"5" :{"A":0.20, "C":0.30, "G":0.30, "T":0.20},
                      "M" :{"A":0.25, "C":0.25, "G":0.25, "T":0.25},
                      "3": {"A":0.35, "C":0.15, "G":0.15, "T":0.35}}
    transition_probs = {"5":{"5": 0.8,"M": 0.2,"3": 0.0},
                        "M":{"5": 0.0,"M": 0.9,"3": 0.1},
                        "3":{"5": 0.0,"M": 0.0,"3": 1.0}}
    
    hmm = HiddenMarkovModel(initial_probs, emission_probs, transition_probs)
    # hmm.set_log(10)
    hmm.view_parameters()

    sequence = "ATGCAATGCGCATGCTAAAA"
    statepath = "555555MMMMMM33333333"
    prob = hmm.joint_score(sequence, statepath)
    print ("Joint score: " + str(prob))

    max_score = -inf
    paths, scores = hmm.viterbi("ACAATGCCGTCTCCGCGACGCCTTTAATTAT")
    print(f"All state path ({len(paths)}):")
    for path,score in zip(paths,scores):
        if score > max_score:
            best_path = path
            max_score = score
        print(path + ": {:0.5f}".format(score))
    print("Optime state path:")
    print(best_path + ": {:0.5f}".format(max_score))

if __name__ == "__main__":
    # test_mat()
    test_HMM()
    pass